# 直接回归Tschebyscheff系数
import argparse
from typing import Dict
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from rsp.antenna.array import linear_array
from apps.wfs.wfs_config import WfsConfig as WG

class TcrMlp(nn.Module):
    def __init__(self, N:int = 8):
        '''
        N为线性阵列天线个数，输入层：天线数量、频率、间距、扫描角度（角度）
        '''
        super().__init__()
        self.l1_fc = nn.Linear(4, 16)
        self.l1_relu = nn.ReLU()
        self.l2_fc = nn.Linear(16, 64)
        self.l2_relu = nn.ReLU()
        self.l3_fc = nn.Linear(64, 8)

    def forward(self, x):
        a1 = self.l1_relu(self.l1_fc(x))
        a2 = self.l2_relu(self.l2_fc(a1))
        return self.l3_fc(a2)

class TcrApp(object):
    def __init__(self):
        self.name = 'apps.wfs.tcr_mlp.TcrMlp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'MLP回归Tschebyscheff系数')
        WG.initialize()
        # TcrMlp.waveform_demo(params=params)
        TcrApp.train_main(params=params)

    @staticmethod
    def train_main(params:Dict = {}) -> None:
        # 载入波束成形相关内容和训练目标
        X = torch.tensor([[WG.N, WG.f, WG.d, WG.theta]], dtype=torch.float32)
        y_label = torch.tensor([[0.26221649, 0.51874705, 0.81196007, 1., 1., 0.81196007, 0.51874705, 0.26221649]], dtype=torch.float32)
        print(f'X: {X.shape}; y_label: {y_label};')
        model = TcrMlp()
        X = X.to(WG.device)
        y_label = y_label.to(WG.device)
        model = model.to(WG.device)
        print(f'The End!')

    @staticmethod
    def waveform_demo(params: Dict = {}) -> None:
        '''
        波束形成总程序
        '''
        N = 8 # 线性天线个数
        f = 300.0E6 # 300MHz频率
        c = 3E8 # 光速
        lambda_ = c / f
        print(f'波长为：lambdam={lambda_};')
        d = 0.5*lambda_ # 天线间距离
        Theta = 5 / 180.0 * torch.pi # 扫描角度，单位为弧度
        x_min, x_max, x_nums = -torch.pi/2.0, torch.pi/2.0, 1000
        thetas = np.linspace(x_min, x_max, x_nums) # 定义绘制扫描角度与天线因子图时横坐标扫描角度的范围
        x_scale = (x_max - x_min) / x_nums
        # 计算天线因子
        # Set up the args
        antenna_type = 'Tschebyscheff'
        side_lobe_level = -30.0
        kwargs = {'number_of_elements': N,
                  'scan_angle': Theta,
                  'element_spacing': d,
                  'frequency': f,
                  'theta': thetas,
                  'window_type': antenna_type,
                  'side_lobe_level': side_lobe_level}
        af = linear_array.array_factor(**kwargs)
        print(f'raw AF: {af.shape}; {af.dtype};')
        y_af = torch.from_numpy(af)
        torch.save(y_af, './work/wfs/y_af.pt')
        bw, sll = linear_array.calculate_bw_sll(20.0 * np.log10(abs(af) + np.finfo(float).eps), x_scale)
        print(f'波束宽度：{bw}; 旁瓣水平：{sll}; N: {N}; Theta: {Theta};')
        fig = plt.figure(1) #Figure()
        axes1 = fig.add_subplot(111)
        axes1.clear()
        axes1.plot(np.degrees(thetas), 20.0 * np.log10(abs(af) + np.finfo(float).eps), '')
        axes1.set_ylim(-80, 5)
        axes1.set_xlabel("Theta (degrees)", size=12)
        axes1.set_ylabel("Array Factor (dB)", size=12)
        axes1.grid(linestyle=':', linewidth=0.5)
        axes1.set_title('Linear Array Antenna Pattern', size=14)
        axes1.tick_params(labelsize=12)
        plt.show()
        print(f'^_^ The End! ^_^')

def main(params:Dict = {}) -> None:
    TcrApp.startup(params=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)